import numpy as np
from PIL import Image
import random


def center_crop_arr(pil_image, image_size):
    """
    Center cropping implementation from ADM.
    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
    """
    while min(*pil_image.size) >= 2 * image_size:
        pil_image = pil_image.resize(
            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
        )

    scale = image_size / min(*pil_image.size)
    pil_image = pil_image.resize(
        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
    )

    arr = np.array(pil_image)
    crop_y = (arr.shape[0] - image_size) // 2
    crop_x = (arr.shape[1] - image_size) // 2
    return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])

def rand_crop_arr(image, image_size):
    image = np.array(image)
    h, w = image.shape[:2]
    side_length = min(h, w)
    x_start = random.randint(0, w - side_length)
    y_start = random.randint(0, h - side_length)
    cropped_image = image[y_start:y_start + side_length, x_start:x_start + side_length]
    return Image.fromarray(cropped_image).resize((image_size, image_size))

def center_rescale_crop_arr(image, image_size):
    image = np.array(image)
    h, w = image.shape[:2]
    side_length = min(h, w)
    x_start = (w - side_length)//2
    y_start = (h - side_length)//2
    cropped_image = image[y_start:y_start + side_length, x_start:x_start + side_length]
    return Image.fromarray(cropped_image).resize((image_size, image_size))
